
#include "efstream.hpp"
#include "sock-inet-udp.hpp"
#include "tftp-packet.hpp"

#include <cstring>
#include <iostream>

#define TIMEOUT 5000000

////////////////////////////////////////////////////////////////////////////////

int main(int argc, const char* argv[])
{
    uint16_t port = 6900;

    if(argc != 2 || sscanf(argv[1], "%hu", &port) != 1)
    {
        std::cout << "using default (6900) port" << std::endl;
    }

    try
    {
        sock_inet_udp server = sock_inet_udp::bind(sock_inet::address(0, port));

        while(true)
        {
            sock_inet_udp::packet_n<tftp_packet::size> req;

            server >> req;

            sock_inet_udp client;

            try
            {
                try
                {
                    tftp_packet_read read(req.get_buffer(), req.get_len());

                    iefstream file(read.get_filename(), std::ifstream::binary);

                    for(size_t sn = 1; !file.eof(); sn++)
                    {
                        client << (tftp_packet_factory(req.get_address()) << tftp_packet::data << sn << file);

                        sock_inet_udp::packet_n<tftp_packet::size> p;

                        while(true)
                        {
                            client >> sock_inet::timeout(TIMEOUT) >> p;

                            if(p.get_address() == req.get_address())
                            {
                                break;
                            }

                            client << (tftp_packet_factory(p.get_address()) << tftp_packet::error << 5 << "Unknown transfer ID");
                        }

                        tftp_packet_ack(p.get_buffer(), p.get_len(), sn);
                    }
                }
                catch(tftp_exception_opcode& e)
                {
                    if(strcmp(e.what(), "invalid read packet"))
                    {
                        throw e;
                    }

                    tftp_packet_write write(req.get_buffer(), req.get_len());

                    oefstream file(write.get_filename(), std::ofstream::binary);

                    for(size_t sn = 0;;)
                    {
                        client << (tftp_packet_factory(req.get_address()) << tftp_packet::ack << sn);

                        sock_inet_udp::packet_n<tftp_packet::size> p;

                        while(true)
                        {
                            client >> sock_inet::timeout(TIMEOUT) >> p;

                            if(p.get_address() == req.get_address())
                            {
                                break;
                            }

                            client << (tftp_packet_factory(p.get_address()) << tftp_packet::error << 5 << "Unknown transfer ID");
                        }

                        tftp_packet_data data(p.get_buffer(), p.get_len(), ++sn);

                        file << data;

                        if(data.get_datalen() < tftp_packet::data_size)
                        {
                            file.close();
                            client << (tftp_packet_factory(req.get_address()) << tftp_packet::ack << sn);
                            break;
                        }
                    }
                }
            }
            catch(exception_timeout&)
            {
                client << (tftp_packet_factory(req.get_address()) << tftp_packet::error << (uint16_t)0 << "Timeout");
            }
            catch(exception_file_not_found&)
            {
                client << (tftp_packet_factory(req.get_address()) << tftp_packet::error << 1 << "File not found");
            }
            catch(exception_access_violation&)
            {
                client << (tftp_packet_factory(req.get_address()) << tftp_packet::error << 2 << "Access violation");
            }
            catch(exception_disk_full&)
            {
                client << (tftp_packet_factory(req.get_address()) << tftp_packet::error << 3 << "Disk full or allocation exceeded");
            }
            catch(tftp_exception_opcode&)
            {
                client << (tftp_packet_factory(req.get_address()) << tftp_packet::error << 4 << "Illegal TFTP operation");
            }
            catch(exception_file_exists&)
            {
                client << (tftp_packet_factory(req.get_address()) << tftp_packet::error << 6 << "File already exists");
            }
            catch(tftp_exception& e)
            {
                client << (tftp_packet_factory(req.get_address()) << tftp_packet::error << (uint16_t)0 << e.what());
            }
            catch(exception_unknown& e)
            {
                client << (tftp_packet_factory(req.get_address()) << tftp_packet::error << (uint16_t)0 << e.what());
            }
        }
    }
    catch(std::exception&)
    {
        perror("tftp");
    }

    return 0;
}

////////////////////////////////////////////////////////////////////////////////
